import os
import argparse
import pandas as pd 


def add_gene_info(cluster_table, pfam_info_file, output_file):
    # 读入基因簇结果表格
    cluster_df = pd.read_excel(cluster_table)
    cluster_title = cluster_df.columns.tolist()
    # print(cluster_title.index("gene_id"))
    # print(cluster_title)
    
    pfam_info = pd.read_csv(pfam_info_file, sep="\t", dtype="str")
    pfam_info_title = pfam_info.columns.tolist() 
    # print(gene_info_title) 
    # print(gene_info)
    pfam_info = pfam_info[pfam_info["Region_type"] != "body"]
    

    # 转录组id作为键，对应行作为值
    gene_info_dict = {}
    for row in pfam_info.values.tolist():
        gene = row[0].strip("transcript:")
        if gene not in gene_info_dict:
            gene_info_dict[gene] = []
        if row[4] not in gene_info_dict[gene]:
            gene_info_dict[gene].append(row[4])
            
    result = []
    for row in cluster_df.values.tolist():
        tmp_row = row 
        gene_name = row[6]
        info = ""
        if gene_name in gene_info_dict:
            info = gene_info_dict[gene_name]
        
        tmp_row = tmp_row + [info]
        result.append(tmp_row)
        
    cluster_title = cluster_title + ["pfam"] 

    result_df = pd.DataFrame(result, columns=cluster_title)
    result_df.to_excel(output_file, index=False)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-cluster","--cluster_table")
    parser.add_argument("-pfam","--pfam_info_file")
    parser.add_argument("-output","--output_file")
    args = parser.parse_args()
    
    return args 


if __name__ == "__main__":
    params = parse_args()
    cluster_table = params.cluster_table
    pfam_info_file = params.pfam_info_file
    output_file = params.output_file
    add_gene_info(cluster_table, pfam_info_file, output_file)